import os
import cv2
import numpy as np
from scipy.spatial import cKDTree
from dataclasses import dataclass
from typing import List, Tuple, Optional
from lane_detection.read_param import CameraTransformation

class CameraTransform:
    """相机变换类。
    
    该类处理相机的标定参数，计算畸变校正和鸟瞰图(BEV)变换。
    
    Attributes:
        mapping_file (str): 映射表文件路径。
        scale_factor (float): 缩放比例。
        cam_matrix (np.ndarray): 相机内参矩阵。
        dist_coeffs (np.ndarray): 畸变系数。
        position (np.ndarray): 棋盘格在相机坐标系中的位置。
        rotation_matrix (np.ndarray): 棋盘格在相机坐标系中的旋转矩阵。
        output_width (int): 输出图像宽度。
        output_height (int): 输出图像高度。
        input_width (int): 输入图像宽度。
        input_height (int): 输入图像高度。
        bev_width (int): BEV图像宽度。
        bev_height (int): BEV图像高度。
        bev_pixels_per_mm (float): BEV每毫米像素数。
        bev_pos (Tuple[float, float]): 相机在BEV中的位置。
        bev_euler (Tuple[float, float, float]): 相机在BEV中的欧拉角。
        M (np.ndarray): 变换矩阵。
        bev_map_x (np.ndarray): BEV映射表X坐标。
        bev_map_y (np.ndarray): BEV映射表Y坐标。
        kd_tree (cKDTree): KD树用于加速点投影。
        indices_to_coords (np.ndarray): 索引到坐标的映射。
        inter_map_x (np.ndarray): 畸变校正映射表X坐标。
        inter_map_y (np.ndarray): 畸变校正映射表Y坐标。
    """
    
    def __init__(self, camera_transformation: CameraTransformation):
        """初始化相机变换类。
        
        Args:
            camera_transformation (CameraTransformation): 相机变换配置参数。
        """
        # 读取标定参数
        calib_params = self.parse_calibration_file(camera_transformation.calib_cam)

        # 检查参数是否完整
        for key in ["camera_matrix", "dist_coeffs", "position", "rotation_matrix"]:
            if calib_params[key] is None:
                print(f"错误：无法从文件中解析 {key}")
                return None, None, None, None
            
        # 存储参数
        self.mapping_file = camera_transformation.mapping_inter                       # 映射表文件路径
        self.scale_factor = camera_transformation.scale_factor                        # 缩放比例
        self.cam_matrix = calib_params["camera_matrix"]                               # 相机内参矩阵
        self.dist_coeffs = calib_params["dist_coeffs"]                                # 畸变系数
        self.position = calib_params["position"]                                      # 棋盘格在相机坐标系中的位置
        self.rotation_matrix = calib_params["rotation_matrix"]                        # 棋盘格在相机坐标系中的旋转矩阵
        self.output_width, self.output_height = camera_transformation.output_shape    # 目标输出图像分辨率 (width, height)
        self.input_width, self.input_height = camera_transformation.input_shape       # 输入图像分辨率 (width, height)
        self.bev_width, self.bev_height = camera_transformation.bev_shape             # BEV图像分辨率 (width, height)
        self.bev_pixels_per_mm = camera_transformation.bev_pixels_per_mm              # BEV图像每毫米对应的像素数
        self.bev_pos = (0, 0)                                                         # 相机在BEV中的位置
        self.bev_euler = (0, 0, 0)                                                    # 相机在BEV中的欧拉角
        self.M = None                                                                 # 变换矩阵
        self.bev_map_x = None                                                         # BEV映射表X
        self.bev_map_y = None                                                         # BEV映射表Y    
        self.kd_tree = None                                                           # KD树
        self.indices_to_coords = None                                                 # 索引到坐标的映射
        self.inter_map_x = None                                                       # 畸变校正映射表X
        self.inter_map_y = None                                                       # 畸变校正映射表Y           
        load_success = False                                                          # 是否成功加载映射表    
        if not camera_transformation.cal_inter_param:
            print("加载预计算的映射表...")
            load_success = self.load_mapping_tables(self.mapping_file)                                                 
        if camera_transformation.cal_inter_param or not load_success:
            # 计算畸变校正映射表
            print("预计算映射表...")
            self.inter_map_x, self.inter_map_y = self.precompute_high_precision_map()
            self.save_mapping_tables(self.mapping_file)

    def parse_calibration_file(self, file_path: str) -> dict:
        """解析标定参数文件。
        
        Args:
            file_path (str): 标定参数文件路径。
            
        Returns:
            dict: 包含相机内参矩阵、畸变系数、位置向量和旋转矩阵的字典。
            
        Raises:
            FileNotFoundError: 如果文件不存在。
            ValueError: 如果文件格式不正确。
        """
        params = {
            "camera_matrix": None,
            "dist_coeffs": None,
            "position": None,
            "rotation_matrix": None
        }
        
        # 检查文件是否存在
        if not os.path.isfile(file_path):
            raise FileNotFoundError(f"标定文件不存在: {file_path}")
        
        # 读取文件内容
        with open(file_path, 'r') as f:
            lines = f.readlines()
            
            # 解析相机内参矩阵
            camera_matrix = []
            for i in range(1, 4):  # 跳过"Camera matrix:"行
                row = lines[i].strip().replace('[', '').replace(']', '').split()
                camera_matrix.append([float(x) for x in row])
            params["camera_matrix"] = np.array(camera_matrix, dtype=np.float64)
            
            # 解析畸变系数
            dist_line = lines[5].strip().replace('[[', '').replace(']]', '').split()
            params["dist_coeffs"] = np.array([float(x) for x in dist_line], dtype=np.float64)
            
            # 解析位置向量
            pos_line1 = lines[7].strip().replace('[', '').replace(']', '').split()
            pos_line2 = lines[8].strip().replace('[', '').replace(']', '').split()
            pos_line3 = lines[9].strip().replace('[', '').replace(']', '').split()
            position = [float(pos_line1[0]), float(pos_line2[0]), float(pos_line3[0])]
            params["position"] = np.array(position, dtype=np.float64).reshape(3, 1)
            
            # 解析旋转矩阵
            rot_matrix = []
            for i in range(11, 14):  # 跳过"Rotation matrix:"行
                row = lines[i].strip().replace('[', '').replace(']', '').split()
                rot_matrix.append([float(x) for x in row])
            params["rotation_matrix"] = np.array(rot_matrix, dtype=np.float64)
        
        return params

    def precompute_high_precision_map(self) -> Tuple[np.ndarray, np.ndarray]:
        """预计算高精度映射表：目标分辨率 -> 原始分辨率。
        
        Returns:
            Tuple[np.ndarray, np.ndarray]: (map_x, map_y) 映射表。
        """
        # 1. 创建目标分辨率的网格
        x_dst, y_dst = np.meshgrid(
            np.arange(self.output_width), 
            np.arange(self.output_height)
        )
        
        # 2. 将目标坐标转换为归一化相机坐标
        # 使用目标分辨率下的等效相机参数
        fx_tgt = self.cam_matrix[0, 0] * (self.output_width / self.input_width)
        fy_tgt = self.cam_matrix[1, 1] * (self.output_height / self.input_height)
        cx_tgt = self.cam_matrix[0, 2] * (self.output_width / self.input_width)
        cy_tgt = self.cam_matrix[1, 2] * (self.output_height / self.input_height)

        # 归一化坐标 (去除内参影响)
        x_norm = (x_dst - cx_tgt) / fx_tgt
        y_norm = (y_dst - cy_tgt) / fy_tgt
        
        # 3. 应用畸变模型
        r2 = x_norm**2 + y_norm**2
        r22 = r2**2
        r23 = r2**3
        radial_dist = 1 + self.dist_coeffs[0]*r2 + self.dist_coeffs[1]*r22 + self.dist_coeffs[4]*r23
        tan_dist_x = 2*self.dist_coeffs[2]*x_norm*y_norm + self.dist_coeffs[3]*(r2 + 2*x_norm**2)
        tan_dist_y = self.dist_coeffs[2]*(r2 + 2*y_norm**2) + 2*self.dist_coeffs[3]*x_norm*y_norm

        # 畸变后的归一化坐标
        x_dist = x_norm * radial_dist + tan_dist_x
        y_dist = y_norm * radial_dist + tan_dist_y
        
        # 4. 映射回原始图像像素坐标
        map_x = x_dist * self.cam_matrix[0, 0] + self.cam_matrix[0, 2]
        map_y = y_dist * self.cam_matrix[1, 1] + self.cam_matrix[1, 2]

        return map_x.astype(np.float32), map_y.astype(np.float32)

    def save_mapping_tables(self, file_path: str):
        """保存映射表到文件。
        
        Args:
            file_path (str): 映射表文件路径。
        """
        np.savez(file_path, map_x=self.inter_map_x, map_y=self.inter_map_y)

    def load_mapping_tables(self, file_path: str) -> bool:
        """从文件加载映射表。
        
        Args:
            file_path (str): 映射表文件路径。
            
        Returns:
            bool: 加载是否成功。
        """
        try:
            data = np.load(file_path)
            self.inter_map_x = data['map_x']
            self.inter_map_y = data['map_y']
            return True
        except Exception as e:
            print(f"加载映射表失败: {str(e)}")
            return False

    def save_bev_mapping_tables(self, file_path: str):
        """保存BEV映射表到文件（优化版）。
        
        Args:
            file_path (str): BEV映射表文件路径。
        """
        save_dict = {
            # 映射表数据
            'bev_map_x': self.bev_map_x.astype(np.float32),  # 减小存储大小
            'bev_map_y': self.bev_map_y.astype(np.float32),
            # KDTree数据（存储点集而非对象）
            'kd_tree_data': self.kd_tree.data if hasattr(self.kd_tree, 'data') else None,
        
            # 坐标映射
            'indices_to_coords': self.indices_to_coords,
            
            # 参数
            'params': {
                'camera_bev_pos': self.bev_pos,
                'camera_bev_euler': self.bev_euler,
                'camera_bev_pos_mm': self.bev_pos_mm,
                'input_shape': (self.input_width, self.input_height),
                'output_shape': (self.output_width, self.output_height),
                'bev_shape': (self.bev_width, self.bev_height),
                'version': '1.0'  # 添加版本控制
            }
        }
    
        # 使用压缩格式存储
        np.savez_compressed(file_path, **save_dict)

    def load_bev_mapping_tables(self, file_path: str) -> bool:
        """从文件加载BEV映射表（优化版）。
        
        Args:
            file_path (str): BEV映射表文件路径。
            
        Returns:
            bool: 加载是否成功。
        """
        try:
            data = np.load(file_path, allow_pickle=True)
            
            # 加载核心数据
            self.bev_map_x = data['bev_map_x']
            self.bev_map_y = data['bev_map_y']
            self.indices_to_coords = data['indices_to_coords']
            
            # 重建KDTree
            kd_tree_data = data['kd_tree_data']
            if kd_tree_data is not None:
                self.kd_tree = cKDTree(kd_tree_data)
            else:
                print("警告: KDTree数据不存在，将重新计算")
                self._build_kd_tree()
            
            # 加载参数
            params = data['params'].item()
            self.bev_pos = params['camera_bev_pos']
            self.bev_euler = params['camera_bev_euler']
            self.bev_pos_mm = params['camera_bev_pos_mm']
            
            # 参数验证
            current_shape = (self.input_width, self.input_height)
            if params['input_shape'] != current_shape:
                print(f"输入尺寸不匹配: 已加载 {params['input_shape']}, 当前 {current_shape}")
                return False
                
            # 版本检查
            if params.get('version', '0.0') != '1.0':
                print("警告: 文件版本不匹配，可能存在兼容性问题")
            print( "加载BEV映射表成功")
            return True
            
        except Exception as e:
            print(f"加载BEV映射表失败: {str(e)}")
            return False
    
    def high_precision_undistort(self, frame: np.ndarray) -> np.ndarray:
        """高精度去畸变：直接从原始图像采样。
        
        Args:
            frame (np.ndarray): 原始图像 (input_width x input_height)。
            
        Returns:
            np.ndarray: 去畸变后的图像 (output_width x output_height)。
            
        Raises:
            ValueError: 如果映射表未初始化。
        """
        if self.inter_map_x is None or self.inter_map_y is None:
            raise ValueError("请先调用 precompute_high_precision_map 计算映射表")
        
        # 应用映射表
        undistorted = cv2.remap(
            frame, 
            self.inter_map_x, 
            self.inter_map_y, 
            interpolation=cv2.INTER_LINEAR  # 使用线性插值保证质量
        )
        return undistorted
    
    def calculate_bev_camera_pose(self, 
                                 bev_center: Tuple[float, float], 
                                 global_rotation: Optional[np.ndarray] = None
                                 ) -> Tuple[Tuple[float, float], Tuple[float, float, float]]:
        """计算相机在BEV空间中的位置和姿态。
        
        Args:
            bev_center (Tuple[float, float]): BEV坐标系的原点在世界坐标系中的位置。
            global_rotation (Optional[np.ndarray]): 全局旋转矩阵（3x3），将世界坐标系旋转到BEV坐标系。
            
        Returns:
            Tuple[Tuple[float, float], Tuple[float, float, float]]: 
                - BEV相机位置 (像素)
                - BEV相机欧拉角 (度)
        """
        # 1. 计算相机在世界坐标系中的位置
        R_inv = self.rotation_matrix.T
        camera_in_world = -R_inv @ self.position

        # 2. 应用全局旋转（如果提供）
        if global_rotation is not None:
            # 将相机位置也从世界坐标系旋转到BEV坐标系
            camera_in_bev = global_rotation @ camera_in_world
        else:
            camera_in_bev = camera_in_world
        
        # 3. 计算相机在BEV空间中的位置
        camera_bev_x = camera_in_bev[0, 0]
        camera_bev_y = camera_in_bev[1, 0]
        
        bev_camera_x = self.bev_width/2 + (camera_bev_x - bev_center[0]) * self.bev_pixels_per_mm
        bev_camera_y = self.bev_height + (camera_bev_y - bev_center[1]) * self.bev_pixels_per_mm
        bev_camera_position = (bev_camera_x, bev_camera_y)
        
        # 4. 计算相机在BEV空间中的欧拉角
        if global_rotation is not None:
            # 相机在BEV坐标系中的姿态 = global_rotation @ R_inv
            R_camera_in_bev = global_rotation @ R_inv
        else:
            R_camera_in_bev = R_inv
        
        # 计算欧拉角
        yaw = np.arctan2(R_camera_in_bev[1, 0], R_camera_in_bev[0, 0])
        pitch = np.arctan2(-R_camera_in_bev[2, 0], np.sqrt(R_camera_in_bev[2, 1]**2 + R_camera_in_bev[2, 2]**2))
        roll = np.arctan2(R_camera_in_bev[2, 1], R_camera_in_bev[2, 2])
        
        bev_camera_euler = np.degrees([roll, pitch, yaw])
        # 存储结果
        self.bev_pos = bev_camera_position
        self.bev_euler = bev_camera_euler
        self.bev_pos_mm = (camera_bev_x, camera_bev_y)

        print("BEV相机位置 (像素):", bev_camera_position)
        print("BEV相机欧拉角 (度):", bev_camera_euler)

        return bev_camera_position, bev_camera_euler
    
    def calculate_perspective_transform(self, 
                                       bev_center: Tuple[float, float], 
                                       global_rotation: Optional[np.ndarray] = None, 
                                       bev_camera_center: Optional[Tuple[float, float]] = None, 
                                       image_yaw_bias: float = 10
                                       ) -> np.ndarray:
        """计算透视变换矩阵。
        
        Args:
            bev_center (Tuple[float, float]): BEV坐标系的原点在世界坐标系中的位置。
            global_rotation (Optional[np.ndarray]): 全局旋转矩阵（3x3），将世界坐标系旋转到BEV坐标系。
            bev_camera_center (Optional[Tuple[float, float]]): 相机在BEV中的中心位置。
            image_yaw_bias (float): 图像yaw旋转偏置（弧度），围绕BEV中心在三维坐标系下的yaw旋转。
            
        Returns:
            np.ndarray: 透视变换矩阵 M。
        """
        # 1. 获取相机高度
        camera_height = self.position[2, 0]

        # 2. 计算视场角 (FOV)
        fov_x = 2 * np.arctan(self.input_width / (2 * self.cam_matrix[0, 0]))
        fov_y = 2 * np.arctan(self.input_height / (2 * self.cam_matrix[1, 1]))

        # 3. 计算可见地面尺寸
        ground_width = 2 * camera_height * np.tan(fov_x / 2)
        ground_height = 2 * camera_height * np.tan(fov_y / 2)

        # 4. 扩展地面范围
        ground_width *= self.scale_factor
        ground_height *= self.scale_factor
        
        # 5. 定义地面上的点（在世界坐标系中）
        world_ground_points = np.array([
            [-ground_width/2, -ground_height/2, 0],  # 左上
            [ ground_width/2, -ground_height/2, 0],  # 右上
            [ ground_width/2,  ground_height/2, 0],  # 右下
            [-ground_width/2,  ground_height/2, 0]   # 左下
        ], dtype=np.float32)
        
        # 6. 应用全局旋转（如果提供）
        if global_rotation is not None:
            # 将世界坐标系中的点旋转到BEV坐标系
            bev_ground_points = []
            for point in world_ground_points:
                bev_point = global_rotation @ point.reshape(3, 1)
                bev_ground_points.append(bev_point.flatten())
            ground_points = np.array(bev_ground_points, dtype=np.float32)
        else:
            # 如果没有全局旋转，使用世界坐标
            ground_points = world_ground_points
                
        # 7. 将世界坐标系中的地面点转换到相机坐标系
        camera_points = []
        for point in world_ground_points:
            # 世界坐标系到相机坐标系的转换
            camera_point = self.rotation_matrix @ point.reshape(3, 1) + self.position
            camera_points.append(camera_point)
        
        # 8. 将相机坐标系投影到图像平面
        image_points = []
        for point in camera_points:
            normalized = point / point[2]
            pixel = self.cam_matrix @ normalized
            image_points.append(pixel[:2])
        
        image_points = np.array(image_points, dtype=np.float32)
        
        # 9. 应用图像yaw旋转偏置（围绕BEV中心在三维坐标系下旋转）
        if image_yaw_bias != 0:
            image_yaw_bias = np.deg2rad(image_yaw_bias)  # 转换为弧度
            
            # 创建yaw旋转矩阵（围绕Z轴）
            cos_yaw = np.cos(image_yaw_bias)
            sin_yaw = np.sin(image_yaw_bias)
            
            # 对地面点应用yaw旋转（围绕BEV中心）
            rotated_ground_points = []
            for point in ground_points:
                # 将点转换为齐次坐标
                point_homogeneous = np.append(point, 1)
                
                # 创建平移矩阵到BEV中心
                translate_to_origin = np.array([
                    [1, 0, 0, -bev_camera_center[0]],
                    [0, 1, 0, -bev_camera_center[1]],
                    [0, 0, 1, 0],
                    [0, 0, 0, 1]
                ], dtype=np.float32)
                
                # 创建旋转矩阵
                rotate_matrix = np.array([
                    [cos_yaw, -sin_yaw, 0, 0],
                    [sin_yaw, cos_yaw, 0, 0],
                    [0, 0, 1, 0],
                    [0, 0, 0, 1]
                ], dtype=np.float32)
                
                # 创建平移矩阵回BEV中心
                translate_back = np.array([
                    [1, 0, 0, bev_camera_center[0]],
                    [0, 1, 0, bev_camera_center[1]],
                    [0, 0, 1, 0],
                    [0, 0, 0, 1]
                ], dtype=np.float32)
                
                # 组合变换：先平移到原点，旋转，再平移回BEV中心
                transform_matrix = translate_back @ rotate_matrix @ translate_to_origin
                
                # 应用变换
                transformed_point = transform_matrix @ point_homogeneous
                
                # 取前三个坐标（去掉齐次坐标的1）
                rotated_ground_points.append(transformed_point[:3])
            
            ground_points = np.array(rotated_ground_points, dtype=np.float32)
        
        # 10. 计算在统一BEV坐标系中的位置
        bev_points = []
        for point in ground_points:
            # 计算相对于BEV原点的偏移
            offset_x = point[0] - bev_center[0]
            offset_y = point[1] - bev_center[1]
            
            # 转换为像素坐标
            bev_x = self.bev_width/2 + offset_x * self.bev_pixels_per_mm
            bev_y = self.bev_height + offset_y * self.bev_pixels_per_mm
            bev_points.append([bev_x, bev_y])
        
        bev_points = np.array(bev_points, dtype=np.float32)
        
        # 11. 计算逆透视变换矩阵
        M = cv2.getPerspectiveTransform(image_points, bev_points)
        self.M = M
        return M
        
    def _build_kd_tree(self):
        """构建KD树加速最近邻搜索。"""
        # 创建点集
        points = np.column_stack((
            self.bev_map_x.ravel(),
            self.bev_map_y.ravel()
        ))
        
        # 构建KD树
        self.kd_tree = cKDTree(points)
        
        # 存储索引到坐标的映射
        self.indices_to_coords = np.array([
            (i % self.bev_width, i // self.bev_width)
            for i in range(len(points))
        ])

    def precompute_bev_maps(self) -> Tuple[np.ndarray, np.ndarray]:
        """预计算从原始图像到BEV的映射表。
        
        Returns:
            Tuple[np.ndarray, np.ndarray]: (bev_map_x, bev_map_y) 映射表。
            
        Raises:
            ValueError: 如果变换矩阵未初始化。
        """
        # 检查是否已经计算了变换矩阵
        if self.M is None:
            raise ValueError("请先调用 calculate_unified_bev_transform_with_rotation 计算变换矩阵")
        
        # 创建BEV图像的网格坐标
        y_bev, x_bev = np.meshgrid(np.arange(self.bev_height), np.arange(self.bev_width), indexing='ij')
        bev_coords = np.stack([x_bev, y_bev, np.ones_like(x_bev)], axis=-1).reshape(-1, 3).T
        
        # 计算BEV坐标对应的原始图像坐标（1080p）
        src_coords_1080p = np.linalg.inv(self.M) @ bev_coords
        src_coords_1080p = src_coords_1080p / src_coords_1080p[2, :]
        src_coords_1080p = src_coords_1080p[:2, :].T.reshape(self.bev_height, self.bev_width, 2)
        
        # 将1080p坐标转换为384x640坐标
        scale_x = self.output_width / self.input_width
        scale_y = self.output_height / self.input_height

        # 计算映射表
        bev_map_x = src_coords_1080p[:, :, 0] * scale_x
        bev_map_y = src_coords_1080p[:, :, 1] * scale_y

        # 转换为float32类型
        self.bev_map_x = bev_map_x.astype(np.float32)
        self.bev_map_y = bev_map_y.astype(np.float32)
        
        # 构建KD树
        self._build_kd_tree()

        return self.bev_map_x, self.bev_map_y

    def fast_bev_transform(self, frame: np.ndarray) -> np.ndarray:
        """快速BEV变换：直接从原始图像采样。
        
        Args:
            frame (np.ndarray): 原始图像 (input_width x input_height)。
            
        Returns:
            np.ndarray: BEV图像 (bev_width x bev_height)。
            
        Raises:
            ValueError: 如果KD树未初始化。
        """
        if self.kd_tree is None:
            raise ValueError("请先调用 precompute_bev_maps 计算映射表")
        
        # 应用映射表
        bev_image = cv2.remap(
            frame,
            self.bev_map_x,
            self.bev_map_y,
            interpolation=cv2.INTER_LINEAR,
        )

        return bev_image
    
    def project_points_with_kdtree(self, lines: List[List[Tuple[float, float]]]) -> List[List[Tuple[float, float]]]:
        """使用KD树加速点投影。
        
        Args:
            lines (List[List[Tuple[float, float]]]): 原始图像中的点列表。
            
        Returns:
            List[List[Tuple[float, float]]]: BEV空间中的点列表。
            
        Raises:
            ValueError: 如果KD树未初始化。
        """
        if self.kd_tree is None:
            raise ValueError("请先调用 precompute_bev_maps 计算映射表")
        
        bev_points = []
        
        # 展平所有点
        all_points = np.array([point for line in lines for point in line])
        
        if len(all_points) == 0:
            return bev_points
        
        # 批量查询最近邻
        _, indices = self.kd_tree.query(all_points)
        
        # 获取BEV坐标
        bev_coords = self.indices_to_coords[indices]
        
        # 重新分组为线
        start = 0
        for line in lines:
            end = start + len(line)
            bev_points.append([tuple(coord) for coord in bev_coords[start:end]])
            start = end
        
        return bev_points